package reactor;

import org.junit.Test;
import org.reactivestreams.Subscription;
import reactor.core.publisher.BaseSubscriber;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import reactor.core.scheduler.Schedulers;

import java.util.concurrent.CountDownLatch;

import static util.Util.print;
import static util.Util.sleep;

/**
 * @author yding5
 */
public class BackPressureTest {

    private static final CountDownLatch countDownLatch = new CountDownLatch(1);

    /**
     * 控制上游流量
     */
    @Test
    public void testRequest() {
        Flux<Integer> flux = Flux.range(1, 10).log();

        flux.subscribe(new BaseSubscriber<>() {
            private int count = 0;
            private final int requestCount = 4;

            @Override
            protected void hookOnSubscribe(Subscription subscription) {
                request(requestCount);
            }

            @Override
            protected void hookOnNext(Integer value) {
                count++;
                // 通过count控制request调用
                if (count == requestCount) {
                    sleep(1000);
                    request(requestCount);
                    count = 0;
                }
            }
        });
    }

    /**
     * 背压策略
     */
    @Test
    public void testBackPressureStrategy() throws InterruptedException {
        // 订阅者
        SlowSubscriber slowSubscriber = new SlowSubscriber();
        // 发布者
        Flux<Integer> fastPublisher =
                createFlux(FluxSink.OverflowStrategy.BUFFER)
                        .doOnRequest(n -> print("下游------>向上游请求" + n + "个数据"))
                        .publishOn(Schedulers.single(), 1);

        fastPublisher.subscribe(slowSubscriber);
        countDownLatch.await();
    }

    private static Flux<Integer> createFlux(FluxSink.OverflowStrategy strategy) {
        return Flux.create(sink -> {
            for (int i = 0; i < 5; i++) {
                sink.next(i);
                print("上游------>数据源创建了数据：" + i);
            }
            sink.complete();
        }, strategy);
    }

    private static class SlowSubscriber extends BaseSubscriber<Integer> {

        @Override
        protected void hookOnSubscribe(Subscription subscription) {
            // 订阅时请求1个数据
            print("订阅开始");
            request(1);
        }

        @Override
        protected void hookOnNext(Integer i) {
            print(("线程" + Thread.currentThread().getName() + "接收数据:" + i));
            request(1);
            sleep(20);
        }

        @Override
        protected void hookOnComplete() {
            print("完成");
            countDownLatch.countDown();
        }
    }
}

